import numpy as np
import pandas as pd
from math import log2


# 计算熵
def calc_entropy(data):
    label_column = data[:, -1]
    unique_labels = np.unique(label_column)
    entropy = 0
    for label in unique_labels:
        prob = np.sum(label_column == label) / len(label_column)
        entropy -= prob * log2(prob)
    return entropy


# 计算信息增益
def calc_information_gain(data, feature_index):
    total_entropy = calc_entropy(data)
    feature_values = np.unique(data[:, feature_index])
    weighted_entropy = 0
    for value in feature_values:
        subset = data[data[:, feature_index] == value]
        subset_entropy = calc_entropy(subset)
        weighted_entropy += (len(subset) / len(data)) * subset_entropy
    information_gain = total_entropy - weighted_entropy
    return information_gain


# 选择最佳的特征
def best_split(data):
    best_feature = -1
    best_info_gain = -1
    for feature_index in range(data.shape[1] - 1):
        info_gain = calc_information_gain(data, feature_index)
        if info_gain > best_info_gain:
            best_info_gain = info_gain
            best_feature = feature_index
    return best_feature


# 创建决策树
def create_tree(data, feature_names):
    labels = data[:, -1]
    unique_labels = np.unique(labels)
    if len(unique_labels) == 1:
        return unique_labels[0]
    if len(feature_names) == 0:
        return np.argmax(np.bincount(labels.astype(int)))
    best_feature = best_split(data)
    best_feature_name = feature_names[best_feature]
    tree = {best_feature_name: {}}
    feature_values = np.unique(data[:, best_feature])
    new_feature_names = [f for i, f in enumerate(feature_names) if i != best_feature]
    for value in feature_values:
        subset = data[data[:, best_feature] == value]
        tree[best_feature_name][value] = create_tree(subset, new_feature_names)
    return tree


# 预测新样本
def predict(tree, sample, feature_names):
    if isinstance(tree, dict):
        feature = list(tree.keys())[0]
        feature_index = feature_names.index(feature)  # 获取特征索引
        value = sample[feature_index]  # 获取样本对应特征的值
        return predict(tree[feature].get(value, -1), sample, feature_names)
    else:
        return tree


# 加载数据集
def load_data(file_path):
    data = pd.read_csv(file_path, header=None)
    return data.values


# 主程序
if __name__ == "__main__":
    file_path = "C:\\Users\\li\\Desktop\\机器学习\\wine\\wine.data"
    data = load_data(file_path)

    feature_names = ['Alcohol', 'Malic_Acid', 'Ash', 'Alcalinity_of_Ash', 'Magnesium', 'Total_Phenols',
                     'Flavanoids', 'Nonflavanoid_Phenols', 'Proanthocyanins', 'Color_Intensity', 'Hue',
                     'OD280_OD315', 'Proline']

    tree = create_tree(data, feature_names)
    print("决策树结构：")
    print(tree)

    sample = data[0, :-1]  # 假设取第一个样本
    prediction = predict(tree, sample, feature_names)
    print(f"预测标签：{prediction}")
